from typing import List, Dict


from conf.config import MYSQL_CONFIG
from db.mysql.mysql_db import MysqlClient


class BuildingInfoTextChunkDAO(object):
    """ 楼盘信息文本块数据访问对象 """

    def __init__(self, *args, **kwargs):
        super(BuildingInfoTextChunkDAO, self).__init__(*args, **kwargs)

    @staticmethod
    def get_mysql_conn():
        """ 获取mysql连接 """
        mysql_conn = MysqlClient(MYSQL_CONFIG)
        return mysql_conn

    def get_text_chunks(self, text_chunk_ids: List[int]) -> List[Dict]:
        """
        获取文本块列表
        :param text_chunk_ids: 文本块的id列表
        :return:
        """
        mysql_conn = self.get_mysql_conn()

        text_chunks = list()
        try:
            sql = ('SELECT text_chunk_id, text_chunk, source_web '
                   'FROM building_info_text_chunk WHERE text_chunk_id in %s '
                   'ORDER BY text_chunk_id;')
            args = (text_chunk_ids, )
            count, text_chunks = mysql_conn.fetchall(sql, args=args)
        finally:
            if "mysql_conn" in dir():  # 判断连接是否成功创建，创建了才能执行close()
                mysql_conn.close()

        return text_chunks


def main():
    building_info_text_chunk_dao = BuildingInfoTextChunkDAO()
    text_chunks = building_info_text_chunk_dao.get_text_chunks(text_chunk_ids=[5, 3, 9])
    print(text_chunks)


if __name__ == '__main__':
    main()
